import math
import numpy as np
from model.search.LocalFeatureSet import LocalFeatureSet
from os.path import join
import os
from model.common.common import *
import logging
import torch
from tqdm import tqdm

def mask_ssc(keypoints, num_ret_points, tolerance, cols, rows):
    
    points_cnt = keypoints.shape[0]
    
    exp1 = rows + cols + 2 * num_ret_points
    exp2 = (
        4 * cols
        + 4 * num_ret_points
        + 4 * rows * num_ret_points
        + rows * rows
        + cols * cols
        - 2 * rows * cols
        + 4 * rows * cols * num_ret_points
    )
    exp3 = math.sqrt(exp2)
    exp4 = num_ret_points - 1

    sol1 = -round(float(exp1 + exp3) / exp4)  # first solution
    sol2 = -round(float(exp1 - exp3) / exp4)  # second solution

    high = (
        sol1 if (sol1 > sol2) else sol2
    )  # binary search range initialization with positive solution
    low = math.floor(math.sqrt(points_cnt / num_ret_points))

    prev_width = -1
    selected_keypoints = []
    result_list = []
    result = []
    complete = False
    k = num_ret_points
    k_min = round(k - (k * tolerance))
    k_max = round(k + (k * tolerance))

    while not complete:
        width = low + (high - low) / 2
        if (
            width == prev_width or low > high
        ):  # needed to reassure the same radius is not repeated again
            result_list = result  # return the keypoints from the previous iteration
            break

        c = width / 2  # initializing Grid
        num_cell_cols = int(math.floor(cols / c))
        num_cell_rows = int(math.floor(rows / c))
        covered_vec = [
            [False for _ in range(num_cell_cols + 1)] for _ in range(num_cell_rows + 1)
        ]
        result = []

        for i in range(points_cnt):
            row = int(
                math.floor(keypoints[i,1] / c)
            )  # get position of the cell current point is located at
            col = int(math.floor(keypoints[i,0] / c))
            if not covered_vec[row][col]:  # if the cell is not covered
                result.append(i)
                # get range which current radius is covering
                row_min = int(
                    (row - math.floor(width / c))
                    if ((row - math.floor(width / c)) >= 0)
                    else 0
                )
                row_max = int(
                    (row + math.floor(width / c))
                    if ((row + math.floor(width / c)) <= num_cell_rows)
                    else num_cell_rows
                )
                col_min = int(
                    (col - math.floor(width / c))
                    if ((col - math.floor(width / c)) >= 0)
                    else 0
                )
                col_max = int(
                    (col + math.floor(width / c))
                    if ((col + math.floor(width / c)) <= num_cell_cols)
                    else num_cell_cols
                )
                for row_to_cover in range(row_min, row_max + 1):
                    for col_to_cover in range(col_min, col_max + 1):
                        if not covered_vec[row_to_cover][col_to_cover]:
                            # cover cells within the square bounding box with width w
                            covered_vec[row_to_cover][col_to_cover] = True

        if k_min <= len(result) <= k_max:  # solution found
            result_list = result
            complete = True
        elif len(result) < k_min:
            high = width - 1  # update binary search range
        else:
            low = width + 1
        prev_width = width
    
    return result_list
    # selected_keypoints = keypoints[result_list,:]
    # # for i in range(len(result_list)):
        
    # #     selected_keypoints.append(keypoints[result_list[i]])

    # return selected_keypoints


class ProcessMaskANMS():
    def __init__(self,context) -> None:
        args = context.args
        self.num_ret_points = 50
        self.tolerance = 0.1
        self.image_size = args.resize
        self.args = args
        self.database = context.dataset

        config = context.config
        self.mask_dir = config["feature_process"]["MaskANMS"]["mask_dir"]


# 00:'road'
# 01:'sidewalk'
# 02:'building'
# 03:'wall'
# 04:'fence'
# 05:'pole'
# 06:'traffic light'
# 07:'traffic sign'
# 08:'vegetation'
# 09:'terrain'
# 10:'sky'
# 11:'person'
# 12:'rider'
# 13:'car'
# 14:'truck'
# 15:'bus'
# 16:'train'
# 17:'motorcycle'
# 18:'bicycle'

        self.mask_type = np.array(
            [1,1,1,1,
             1,1,1,1,
             0,1,0,0,
             0,0,0,0,
             0,0,0])
        pass
    
    def run(self,in_featureset,out_featureset):
        logging.debug("START Mask ANMS")
        # f_set2 = LocalFeatureSet(self.args)
        # dataset_name = self.args.dataset_name
        # f_set2._feature_dir = join(self.args.datasets_folder, dataset_name, "images","test","features_MaskANMS_" + timestr())
        # os.mkdir(f_set2._feature_dir)
        # result={}
        
        featureset= in_featureset.get_features()
        count = len(featureset.items())
        pbar = tqdm(total=count)
        for imgid,features in featureset.items():
            pbar.update(1)
            keypoints = features["keypoints"]

            filename = self.database.get_filename_byid(imgid)
            filename = os.path.basename(filename)
            maskfile_path = self.mask_dir + "/" + filename + ".pt"
            mask_tensor =  torch.load(maskfile_path)
            mask_nparr = mask_tensor.cpu().numpy()

            points_cnt = keypoints.shape[0]

            id_list = []
            for i in range(points_cnt):
                pt = keypoints[i]
                col = int(math.floor(pt[0]))#x
                row = int(math.floor(pt[1]))#y

                #行在前，列在后
                mv = mask_nparr[row,col]
                if self.mask_type[mv] > 0:
                    id_list.append(i)
         
            keypoints2 = keypoints[id_list,:]
            descriptors2 = features["descriptors"][id_list,:]

            #keypoints必须按照响应值倒序排序
            #暂时没有进行排序
            # x_index = np.lexsort(keypoints[:,2])
            # x_index = x_index[::-1]
            # keypoints = keypoints[x_index]
            # result_list = mask_ssc(keypoints2,self.num_ret_points,self.tolerance,self.image_size[1],self.image_size[0])
            # selected_keypoints = keypoints2[result_list,:]
            # selected_descriptors = descriptors2[result_list,:]
            
            selected_keypoints = keypoints2
            selected_descriptors = descriptors2
            
            featuresnew = {}
            featuresnew["keypoints"] = selected_keypoints
            featuresnew["descriptors"] = selected_descriptors
            
            out_featureset.add_features(imgid,featuresnew)
        out_featureset.write()
        pbar.close()
        return 0

    def run_img(self,img):
            features = img["features"]
            imgid = img["id"]

            keypoints = features["keypoints"]

            filename = self.database.get_filename_byid(imgid)
            filename = os.path.basename(filename)
            
            maskfile_path = self.mask_dir + "/" + filename + ".pt"

            mask_tensor =  torch.load(maskfile_path)
            mask_nparr = mask_tensor.cpu().numpy()

            points_cnt = keypoints.shape[0]

            id_list = []
            for i in range(points_cnt):
                pt = keypoints[i]
                y = int(math.floor(pt[1]))
                x = int(math.floor(pt[0]))

                mv = mask_nparr[y,x]
                if self.mask_type[mv] > 0:
                    id_list.append(i)
         
            keypoints2 = keypoints[id_list,:]
            descriptors2 = features["descriptors"][id_list,:]

            result_list = mask_ssc(keypoints2,self.num_ret_points,self.tolerance,self.image_size[1],self.image_size[0])
            
            selected_keypoints = keypoints2[result_list,:]
            selected_descriptors = features["descriptors"][result_list,:]
            
            featuresnew = {}
            featuresnew["keypoints"] = selected_keypoints
            featuresnew["descriptors"] = selected_descriptors

            return featuresnew